import random,time

from paho.mqtt import client as mqtt_c
from locust import TaskSet,task,User,constant_pacing
import redis
pool=redis.ConnectionPool(host='127.0.0.1',port=6379,db=0)
r = redis.StrictRedis(connection_pool=pool)

# broker_add = "112.124.54.250"
broker_add = "127.0.0.1"
port = 1883
topic = "/python/mqtt_topic_for_python"


class TheTaskSet(TaskSet):
    client = mqtt_c.Client()    

    def __init__(self, *args, **kwargs):
        super(TheTaskSet, self).__init__(*args, **kwargs)

        def on_connect(client,userdata,flags,rc):
            if rc == 0:
                print("Connected to MQTT Broker!")
            else:
                print("Failed to connect, return code %d \n",rc)
        
        client_id = f"python-mqtt-{r.incr('count')}"
        print("创建的时候执行 %s",client_id)
        self.client = mqtt_c.Client(client_id)
        self.client.on_connect= on_connect

    def on_start(self):
        print("开始的时候执行")
        self.client.connect(broker_add,port)
        self.client.loop_start()

    def on_stop(self):
        print("结束的时候执行")
    @task
    def task_1(self):
        print("运行的时候执行")
        self.publish()

    def publish(self):
        time.sleep(1)
        msg = f"message: 0"
        result = self.client.publish(topic,msg)
        status = result[0]
        if status == 0:
            print(f"send `{msg}` to topic `{topic}` ")
        else:
            print(f"Failed to send message to topic {topic}")


class TheUser(User):
    tasks = [TheTaskSet]
    wait_time = constant_pacing(1)
